"""
ERNIE-4.5-0.3B-Paddle模型封装类
提供智能文本生成和任务分析功能
"""

import os
import json
import numpy as np
from typing import List, Dict, Any, Optional
import logging

try:
    import paddle
    import paddlenlp
    from paddlenlp.transformers import AutoTokenizer, AutoModelForCausalLM
except ImportError:
    print("警告: PaddlePaddle或PaddleNLP未安装，请运行: pip install paddlepaddle paddlenlp")

from config import MODEL_CONFIG, ZHIYUAN_CONFIG

class ERNIETaskPlanner:
    """
    基于ERNIE-4.5-0.3B-Paddle的智能任务规划器
    结合智源研究院数据进行深度分析
    """
    
    def __init__(self, model_path: str = None):
        self.model_path = model_path or MODEL_CONFIG["model_path"]
        self.tokenizer = None
        self.model = None
        self.knowledge_base = self._load_zhiyuan_knowledge()
        self._initialize_model()
    
    def _initialize_model(self):
        """初始化ERNIE模型"""
        try:
            if os.path.exists(self.model_path):
                print(f"正在加载ERNIE模型: {self.model_path}")
                self.tokenizer = AutoTokenizer.from_pretrained(self.model_path)
                self.model = AutoModelForCausalLM.from_pretrained(self.model_path)
                print("ERNIE模型加载成功!")
            else:
                print(f"模型路径不存在: {self.model_path}")
                print("使用模拟模式运行...")
                self._use_simulation_mode()
        except Exception as e:
            print(f"模型加载失败: {e}")
            print("使用模拟模式运行...")
            self._use_simulation_mode()
    
    def _use_simulation_mode(self):
        """模拟模式，用于演示"""
        self.simulation_mode = True
        print("启用模拟模式 - 将使用预设的智能响应")
    
    def _load_zhiyuan_knowledge(self) -> Dict[str, Any]:
        """加载智源研究院相关知识"""
        return {
            "productivity_principles": [
                "艾森豪威尔矩阵：根据紧急性和重要性划分任务",
                "番茄工作法：25分钟专注工作+5分钟休息",
                "GTD方法：收集、处理、组织、回顾、执行",
                "时间块管理：为不同类型任务分配专门时间段",
                "能量管理：在精力最佳时处理最重要任务"
            ],
            "cognitive_patterns": [
                "人类注意力周期：90-120分钟的自然节律",
                "认知负荷理论：避免同时处理过多复杂任务",
                "心流状态：深度专注的最佳工作状态",
                "决策疲劳：避免在疲劳时做重要决策"
            ],
            "optimization_strategies": [
                "批处理相似任务提高效率",
                "利用生物钟安排高认知需求任务",
                "设置缓冲时间应对突发情况",
                "定期回顾和调整计划"
            ]
        }
    
    def analyze_task_priority(self, task: str, context: str = "") -> Dict[str, Any]:
        """
        分析任务优先级
        结合智源研究院的认知科学数据
        """
        prompt = f"""
基于智源研究院的生产力管理研究和认知科学理论，分析以下任务的优先级：

任务：{task}
上下文：{context}

请从以下维度分析：
1. 紧急性（1-10分）
2. 重要性（1-10分）
3. 复杂度（1-10分）
4. 预估时间（分钟）
5. 最佳执行时间段
6. 所需认知资源级别

请以JSON格式返回分析结果。
"""
        
        if hasattr(self, 'simulation_mode'):
            return self._simulate_task_analysis(task)
        
        return self._generate_response(prompt, "task_analysis")
    
    def _simulate_task_analysis(self, task: str) -> Dict[str, Any]:
        """模拟任务分析结果"""
        # 基于任务关键词的简单分析
        urgency = 7 if any(word in task.lower() for word in ['紧急', '立即', '今天', '截止']) else 5
        importance = 8 if any(word in task.lower() for word in ['重要', '关键', '核心', '项目']) else 6
        complexity = 7 if any(word in task.lower() for word in ['复杂', '研究', '分析', '设计']) else 5
        
        return {
            "urgency": urgency,
            "importance": importance,
            "complexity": complexity,
            "estimated_time": complexity * 15,  # 简单估算
            "optimal_time_slot": "上午" if importance > 7 else "下午",
            "cognitive_load": "高" if complexity > 6 else "中",
            "priority_score": urgency + importance,
            "category": self._categorize_task(task)
        }
    
    def _categorize_task(self, task: str) -> str:
        """任务分类"""
        if any(word in task.lower() for word in ['会议', '沟通', '讨论']):
            return "沟通协作"
        elif any(word in task.lower() for word in ['写', '文档', '报告']):
            return "文档创作"
        elif any(word in task.lower() for word in ['分析', '研究', '学习']):
            return "分析思考"
        elif any(word in task.lower() for word in ['开发', '编程', '代码']):
            return "技术开发"
        else:
            return "其他任务"
    
    def generate_schedule(self, tasks: List[Dict], preferences: Dict = None) -> Dict[str, Any]:
        """
        生成智能作息规划表
        基于智源研究院的时间管理和认知优化研究
        """
        prompt = f"""
基于智源研究院的时间管理研究和认知科学理论，为以下任务生成最优的作息规划表：

任务列表：{json.dumps(tasks, ensure_ascii=False, indent=2)}
用户偏好：{json.dumps(preferences or {}, ensure_ascii=False, indent=2)}

请考虑以下因素：
1. 人类认知节律和注意力周期
2. 任务的认知负荷匹配最佳时间段
3. 艾森豪威尔矩阵优先级原则
4. 番茄工作法和时间块管理
5. 适当的休息和缓冲时间

请生成详细的时间安排表，包括：
- 具体时间段
- 任务安排
- 休息时间
- 优化建议

以JSON格式返回结果。
"""
        
        if hasattr(self, 'simulation_mode'):
            return self._simulate_schedule_generation(tasks)
        
        return self._generate_response(prompt, "schedule_generation")
    
    def _simulate_schedule_generation(self, tasks: List[Dict]) -> Dict[str, Any]:
        """模拟日程生成"""
        # 按优先级排序任务
        sorted_tasks = sorted(tasks, key=lambda x: x.get('priority_score', 0), reverse=True)
        
        schedule = {
            "date": "今天",
            "time_blocks": [],
            "total_work_time": 0,
            "break_time": 0,
            "optimization_tips": []
        }
        
        current_time = 9 * 60  # 9:00 AM in minutes
        
        for task in sorted_tasks:
            duration = task.get('estimated_time', 30)
            
            # 添加任务时间块
            start_time = f"{current_time // 60:02d}:{current_time % 60:02d}"
            end_time = f"{(current_time + duration) // 60:02d}:{(current_time + duration) % 60:02d}"
            
            schedule["time_blocks"].append({
                "start_time": start_time,
                "end_time": end_time,
                "task": task.get('task', '未知任务'),
                "type": "工作",
                "priority": task.get('priority_score', 0),
                "category": task.get('category', '其他')
            })
            
            current_time += duration
            schedule["total_work_time"] += duration
            
            # 添加休息时间
            if duration > 45:  # 长任务后休息15分钟
                break_duration = 15
                break_start = f"{current_time // 60:02d}:{current_time % 60:02d}"
                current_time += break_duration
                break_end = f"{current_time // 60:02d}:{current_time % 60:02d}"
                
                schedule["time_blocks"].append({
                    "start_time": break_start,
                    "end_time": break_end,
                    "task": "休息时间",
                    "type": "休息",
                    "priority": 0,
                    "category": "休息"
                })
                
                schedule["break_time"] += break_duration
        
        # 添加优化建议
        schedule["optimization_tips"] = [
            "高优先级任务安排在上午精力最佳时段",
            "每45-90分钟安排适当休息",
            "相似类型任务批量处理提高效率",
            "预留20%缓冲时间应对突发情况"
        ]
        
        return schedule
    
    def _generate_response(self, prompt: str, task_type: str) -> Dict[str, Any]:
        """生成模型响应"""
        try:
            inputs = self.tokenizer(prompt, return_tensors="pd", max_length=MODEL_CONFIG["max_length"], truncation=True)
            
            with paddle.no_grad():
                outputs = self.model.generate(
                    inputs["input_ids"],
                    max_length=MODEL_CONFIG["max_length"],
                    temperature=MODEL_CONFIG["temperature"],
                    top_p=MODEL_CONFIG["top_p"],
                    repetition_penalty=MODEL_CONFIG["repetition_penalty"],
                    do_sample=True
                )
            
            response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
            
            # 尝试解析JSON响应
            try:
                # 提取JSON部分
                json_start = response.find('{')
                json_end = response.rfind('}') + 1
                if json_start != -1 and json_end != -1:
                    json_str = response[json_start:json_end]
                    return json.loads(json_str)
            except:
                pass
            
            return {"raw_response": response, "type": task_type}
            
        except Exception as e:
            print(f"生成响应时出错: {e}")
            return {"error": str(e), "type": task_type}
    
    def get_productivity_insights(self) -> List[str]:
        """获取基于智源研究院数据的生产力洞察"""
        return [
            "🧠 认知科学研究表明：上午9-11点是大脑最活跃的时段，适合处理复杂任务",
            "⏰ 时间管理专家建议：使用时间块技术，为不同类型任务分配专门时段",
            "🎯 艾森豪威尔矩阵：优先处理重要且紧急的任务，避免陷入忙碌陷阱",
            "🍅 番茄工作法：25分钟专注工作+5分钟休息，提高注意力和效率",
            "🔄 GTD方法论：定期回顾和更新任务列表，保持系统的有效性",
            "⚡ 能量管理：在精力最佳时处理最重要的工作，避免决策疲劳",
            "🧘 认知负荷理论：避免同时处理多个复杂任务，专注单一目标"
        ]